% Estimate dynamic factor model as described in Ritschl, Sarferaz, and Uebele (2014)

clear all
clc

addpath('../funcs');
addpath('../scrpts');
addpath('../data');
%% Prior for factor loadings
diffuse = [12 0.01];
intermediate = [102 0.1];
tight = [1002 0.001];
load('../matFiles/prior_facload');
%% Load and prepare data
data52= load('data/US_53_1867_2006_version_8_2009.txt');
data98= load('US_98_1867_1939_version_10_2008.txt');

subset = 'nonagr_real'; % choose subset
large=0; % choose dataset (for large=1 set subset = 'data98')

if large==1
    data=data98;
else
    data = data52;
end
% Apply HP filter
data_hp = dataPrep(data,6.25,15,2,3);
% Define subsets
setsubset_new;

if size(data_hp,2)>53
    yt = data_hp';
    
    filename = char(['../matFiles/',prior,'/results_US_98_1867_1929_',prior]);
else
    yt = data_hp(:,eval(subset))';
    
    filename = char(['../matFiles/',prior,'/results_52_1867_2006_']);
    filename=char([filename,subset,'_',prior]);
end


%% Gibbs sampling preliminaries
chain           = 100000;   % number of draws for gibbs sampler
burn            =  chain*0.8; % burn-in phase
jumpsave        = ceil((chain-burn)/1000);    % number of draws that are skipped
chain_save      = (chain-burn)/jumpsave;
gx              = 1;

%% Model specifications
K = 1; % number of factors
q=8; % lag length of factors
p=1; % lag length of idiosyncratic components
m = K*q;

stat = 1; % stationarity check (stat=1) and without (stat=0)
sv = 0; % stochastic volatlity sampler: Watanabe and Omori (sv=1) and Kim et al. (sv=0)
trunc = 0; % truncate first factor loading to be positive (trunc=1)

qs     = [0.0073;0.10556;0.00002;0.04395;0.34001;0.24566;0.2575];
mu     = [-10.12999;-3.97281;-8.56686;2.77786;0.61942;1.79518;-1.08819];
v2     = [5.79596;2.61369;5.17950;0.16735;0.64009;0.34023;1.26261];


%% Number of observations
[N,T] = size(yt);
n = K*(p+1);
%T = T+1;
%% Priors

% Prior for the VAR parameters
tau_f = 0.2; % uncertainty around prior mean of zero of factor coeff.
tau_u = 1; % uncertainty around prior mean of zero of idio. coeff.

% Prior for (V)AR coefficients

% Zero restrictions
VPhi_prior = ones((K)*q,K)*1e-12;

% Prior for those parameters without zero restrictions

% Prior for Phi
for ix=1:K
    for qx=1:q
        VPhi_prior(K*(qx-1)+ix,ix) = tau_f/qx;
    end
end

% Prior for Theta
%for ix=1:N
for px=1:p
    VTheta_prior((px-1)+ix,ix) = tau_u/px;
end
%end


Phi_prior = zeros((K)^2*q,1);
VPhi_prior = diag(vec(VPhi_prior));

Theta_prior = zeros(p,1);
VTheta_prior = diag(vec(VTheta_prior));

   
if (strcmp(subset,'agr')==1 ||strcmp(subset,'nonagr_real')==1) % agr and non_agr tend to be unstable --> somewhat tighter prior
    priorlam = [72,0.01];
else
    priorlam = eval(prior);
end

Tlam_0= priorlam(1);% prior for shape of VCV
Qlam_0=eye(N*K)*priorlam(2);% prior for scale of VCV

Tu_0= 6;% prior for shape of VCV
Qu_0=eye(N)*0.001;% prior for scale of VCV

priorvol = [102 0.1];%eval(prior);
Th_0= priorvol(1);% prior for shape of VCV
Qh_0=eye(K)*priorvol(2);% prior for scale of VCV

%% Initial conditions
f0 = zeros(m,1);
Pf0= eye(m);

lam0 = zeros(n,1);
Plam0= eye(n);

h0 = zeros(K,1);
Ph0= zeros(K);
%% Starting Values
ft_pc = extract(yt')';
[maxcorr,indx] = max(corr(ft_pc',yt'));

ft    = ft_pc;

%lamt   = randn(N,K,T);
htq=randn(K,T-q)*0.001;
Phi = VAROLS(ft',q,1);
Theta = eye(N)*0.1;
Omega_chi = eye(N)*0.1;
Omega_eps = eye(N)*0.1;
Omega_xi = eye(K,1)*0.1;

sign=0;
wsign = 0;
wstato = 0;

%% Pre-allocate Space

results.ft          = zeros(K,T,chain_save);
results.lamt        = zeros(N,K,T,chain_save);
results.ht          = zeros(K,T-q,chain_save);
results.Phi         = zeros(K,K*q,chain_save);
results.Theta       = zeros(N,N*p,chain_save);
results.Omega_chi   = zeros(N,N,chain_save);
results.Omega_eps   = zeros(N,N,chain_save);
results.Omega_xi    = zeros(K,chain_save);
results.nut         = zeros(N,T-1,chain_save);

results.stato_f     = zeros(1,chain_save);
results.wsign       = zeros(1,chain_save);

lamvec = zeros(N*K,T);
tic
for w=1:chain
    
    if mod(w,100)==0
        [ssff errorm]   = sprintf('Sampler at %d out of %d draws, subset: %s',w,chain,subset);
        disp(ssff);
    end
    
    %% Quasi-difference data
    yt_star = quasiDiff(yt,T,N,p,Theta);
    
    %% Draw factor loadings
    for nx=1:N
        if trunc==1
            if nx==indx
                [lamvec(nx,:),sign] = filterSmootherCK_trunc(yt_star(nx,:),T,1,n,1,p+1,lam0,Plam0,@state_space_lam,Omega_eps(nx,nx),Omega_chi(nx,nx),eye(1),Theta(nx,nx),ft);
            else
                lamvec(nx,:) = filterSmootherCK(yt_star(nx,:),T,1,n,1,p+1,lam0,Plam0,@state_space_lam,Omega_eps(nx,nx),Omega_chi(nx,nx),eye(1),Theta(nx,nx),ft);
            end
        else
            lamvec(nx,:) = filterSmootherCK(yt_star(nx,:),T,1,n,1,p+1,lam0,Plam0,@state_space_lam,Omega_eps(nx,nx),Omega_chi(nx,nx),eye(1),Theta(nx,nx),ft);
        end
    end
    
    if sign==1
        wsign = wsign +1;
    end
    
    lamt = reshape(lamvec,[N,K,T]);
    
    %% Draw factors
    ht = [zeros(K,q) htq];
    if sv==1
        ft = filterSmootherCK(yt_star,T,N,m,K,q,f0,Pf0,@state_space_f,exp(ht),Omega_chi,Phi,Theta,q,p,lamt);
    else
        ft = filterSmootherCK(yt_star,T,N,m,K,q,f0,Pf0,@state_space_f,exp(ht).^2,Omega_chi,Phi,Theta,q,p,lamt);
    end
    
    %% Draw (V)AR coefficient of factors
    [Phi, stato_f, nu_star] = alphaVAR(ft',htq,K,q,VPhi_prior,Phi_prior,stat,sv);
    
    if stato_f == 1
        wstato = wstato + 1;
    end
    
    %% Draw (V)AR coefficient of idionsyncratic components
    ut = zeros(N,T-1);
    for tx=1:T-1
        ut(:,tx) = yt(:,tx)-squeeze(lamt(:,:,tx))*ft(:,tx);
    end
    
    for nx=1:N
        Theta(nx,nx) = alphaVAR(ut(nx,:)',Omega_chi(nx,nx),1,p,VTheta_prior,Theta_prior,stat,2);
    end
    
    %% Draw stochastic volatilities
    if sv==1
        for ix=1:K
            htq(ix,:) = (svsamp(nu_star(:,ix),htq(ix,:)',Omega_xi(ix), h0(ix), Ph0(ix,ix),3))';
        end
    else
        % Compute st
        st=gen_s(T-q,K,nu_star,ht(:,q+1:end),qs,mu,v2);
        %Draw stochastic volatilities
        htq = filterSmootherCK(nu_star,T-q,K,K,K,1,h0,Ph0,@state_space_h,st,mu,v2,Omega_xi);
    end
    
    %% Draw variance of idionsyncratic equation
    Omega_chi = alphaVariance(ut,Theta,eye(N),p,Qu_0,Tu_0);
    
    %% Draw variance of factor loadings
    [Omega_eps,nut] = alphaVariance(lamvec,eye(N*K),eye(K),1,Qlam_0,Tlam_0);
    %% Draw variance of stochastic volatilities
    Omega_xi = alphaVariance(htq,eye(K),eye(K),1,Qh_0,Th_0);
    %% Save draws
    
    if w>burn
        if mod(w,jumpsave)==0
            results.ft(:,:,gx) = ft;
            results.lamt(:,:,:,gx) = lamt;
            results.ht(:,:,gx) = htq;
            results.Phi(:,:,gx) = Phi;
            results.Theta(:,:,gx) = Theta;
            results.Omega_chi(:,:,gx) = Omega_chi;
            results.Omega_eps(:,:,gx) = Omega_eps;
            results.Omega_xi(:,gx) = Omega_xi;
            results.nut(:,:,gx) = nut;
            
            results.stato_f(:,gx) = stato_f;
            results.wsign(:,gx)   = sign;
            
            gx                      = gx + 1;
        end
    end
    
end

toc
save(filename ,'results', 'yt','T','q','N','K','Phi_prior','VPhi_prior',...
    'Theta_prior', 'VTheta_prior', 'Qlam_0', 'Tlam_0', 'Qu_0', 'Tu_0','Qh_0', 'Th_0',...
    'f0', 'Pf0', 'h0', 'Ph0', 'lam0', 'Plam0', 'sv', 'ft_pc', 'chain', 'chain_save');
